import socket
import time
import threading
import queue
from enum import Enum, auto


class ConnectionState(Enum):
    DISCONNECTED = auto()
    CONNECTING = auto()
    CONNECTED = auto()
    RECONNECTING = auto()


class EnhancedTCPClient:
    def __init__(self, host='127.0.0.1', port=65432):
        self.host = host
        self.port = port
        self.client_socket = None
        self.state = ConnectionState.DISCONNECTED
        self.reconnect_attempts = 0
        self.max_reconnect_attempts = 5
        self.reconnect_delay = 1  # initial delay in seconds
        self.max_reconnect_delay = 30  # max delay in seconds
        self.heartbeat_interval = 10  # seconds
        self.last_heartbeat = 0
        self.running = False
        self.lock = threading.Lock()

    def connect(self):
        with self.lock:
            if self.state in [ConnectionState.CONNECTED, ConnectionState.CONNECTING]:
                return True

            self.state = ConnectionState.CONNECTING
            print("Attempting to connect...")

        try:
            self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            # Enable TCP keepalive
            self.client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            # Linux specific: TCP keepalive time (seconds since last data sent)
            self.client_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60)
            # Linux specific: Interval between keepalive probes
            self.client_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 10)
            # Linux specific: Number of unacknowledged probes before declaring connection dead
            self.client_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 3)

            self.client_socket.connect((self.host, self.port))

            with self.lock:
                self.state = ConnectionState.CONNECTED
                self.reconnect_attempts = 0
                self.reconnect_delay = 1
                self.last_heartbeat = time.time()

            print(f"Connected to server {self.host}:{self.port}")
            return True

        except Exception as e:
            with self.lock:
                self.state = ConnectionState.DISCONNECTED
            print(f"Connection failed: {e}")
            return False

    def auto_reconnect(self):
        with self.lock:
            if self.state == ConnectionState.RECONNECTING:
                return
            self.state = ConnectionState.RECONNECTING

        print("Starting auto-reconnect...")

        while self.running and self.reconnect_attempts < self.max_reconnect_attempts:
            time.sleep(self.reconnect_delay)

            if self.connect():
                return

            with self.lock:
                self.reconnect_attempts += 1
                self.reconnect_delay = min(
                    self.reconnect_delay * 2,  # exponential backoff
                    self.max_reconnect_delay
                )
                print(f"Reconnect attempt {self.reconnect_attempts}, next in {self.reconnect_delay}s")

        with self.lock:
            if self.state != ConnectionState.CONNECTED:
                self.state = ConnectionState.DISCONNECTED
                print("Max reconnect attempts reached. Giving up.")

    def send_heartbeat(self):
        if not self.running or self.state != ConnectionState.CONNECTED:
            return

        current_time = time.time()
        if current_time - self.last_heartbeat >= self.heartbeat_interval:
            try:
                self.send_data("HEARTBEAT")
                with self.lock:
                    self.last_heartbeat = current_time
            except Exception as e:
                print(f"Heartbeat failed: {e}")
                self.handle_disconnection()

    def send_data(self, data):
        if not self.running or self.state != ConnectionState.CONNECTED:
            raise Exception("Not connected to server")

        try:
            self.client_socket.sendall(data.encode())
            response = self.client_socket.recv(1024)
            print(f"Server response: {response.decode()}")
            return True
        except Exception as e:
            print(f"Send/receive error: {e}")
            self.handle_disconnection()
            raise

    def handle_disconnection(self):
        with self.lock:
            if self.state == ConnectionState.DISCONNECTED:
                return

            self.state = ConnectionState.DISCONNECTED
            if self.client_socket:
                self.client_socket.close()

        if self.running:
            reconnect_thread = threading.Thread(target=self.auto_reconnect)
            reconnect_thread.daemon = True
            reconnect_thread.start()

    def start(self):
        self.running = True
        if not self.connect():
            self.auto_reconnect()

        # Start heartbeat thread
        threading.Thread(target=self.heartbeat_loop, daemon=True).start()

    def heartbeat_loop(self):
        while self.running:
            self.send_heartbeat()
            time.sleep(1)

    def stop(self):
        self.running = False
        with self.lock:
            if self.client_socket:
                self.client_socket.close()
        self.state = ConnectionState.DISCONNECTED
        print("Client stopped")


if __name__ == "__main__":
    client = EnhancedTCPClient()
    client.start()

    try:
        while True:
            if client.state == ConnectionState.CONNECTED:
                try:
                    message = input("Enter message to send (or 'quit' to exit): ")
                    if message.lower() == 'quit':
                        break
                    client.send_data(message)
                except Exception as e:
                    print(f"Error: {e}")
            else:
                print("Waiting for connection...")
                time.sleep(1)
    finally:
        client.stop()


class ScheduledTCPClient(EnhancedTCPClient):
    """实现定时发送"""
    def __init__(self, host='127.0.0.1', port=65432, send_interval=5):
        super().__init__(host, port)
        self.send_interval = send_interval
        self.message_counter = 0

    def start_scheduled_sending(self):
        def sending_loop():
            while self.running:
                if self.state == ConnectionState.CONNECTED:
                    self.message_counter += 1
                    message = f"Scheduled message {self.message_counter}"
                    try:
                        self.send_data(message)
                    except Exception as e:
                        print(f"Scheduled send failed: {e}")
                time.sleep(self.send_interval)

        threading.Thread(target=sending_loop, daemon=True).start()


class QueuedTCPClient(EnhancedTCPClient):
    """消息队列与断线缓存"""
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.message_queue = queue.Queue()
        self.max_queue_size = 100

    def send_data(self, data):
        if self.state != ConnectionState.CONNECTED:
            if self.message_queue.qsize() < self.max_queue_size:
                self.message_queue.put(data)
                print(f"Message queued (total: {self.message_queue.qsize()})")
            else:
                print("Queue full, message discarded")
            raise Exception("Not connected")

        try:
            super().send_data(data)
            # Send queued messages if any
            while not self.message_queue.empty():
                queued_msg = self.message_queue.get()
                super().send_data(queued_msg)
                print(f"Sent queued message: {queued_msg}")
        except Exception as e:
            raise


class MultiServerTCPClient(EnhancedTCPClient):
    """多服务器故障转移"""
    def __init__(self, servers, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.servers = servers  # list of (host, port) tuples
        self.current_server_index = 0

    def connect(self):
        for i in range(len(self.servers)):
            self.host, self.port = self.servers[self.current_server_index]
            if super().connect():
                return True

            # Try next server
            self.current_server_index = (self.current_server_index + 1) % len(self.servers)

        return Falseclass MultiServerTCPClient(EnhancedTCPClient):
    def __init__(self, servers, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.servers = servers  # list of (host, port) tuples
        self.current_server_index = 0

    def connect(self):
        for i in range(len(self.servers)):
            self.host, self.port = self.servers[self.current_server_index]
            if super().connect():
                return True

            # Try next server
            self.current_server_index = (self.current_server_index + 1) % len(self.servers)

        return False


if __name__ == "__main__":
    client = ScheduledTCPClient(send_interval=3)
    client.start()
    client.start_scheduled_sending()

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        client.stop()